function [psi,phi] = cmr(psi,phi,Q,in,mu,w,dx,g)
% function [psi,phi] = cmr(psi,phi,Q,in,mu,w,dx)
%   This function performs one iteration of coarse mesh rebalance,
%   following the formulation given by Lewis and Miller.
%
%   Inputs:
%           psi     -- fine mesh angular flux
%           phi     -- fine mesh scalar flux
%           Q       -- fine mesh external source
%           in      -- inpute structure
%           mtt     -- material index for each fine mesh
%           mu      -- angle set
%           w       -- angle weight
%           dx      -- fine mesh delta x
%  
%   Example discretization and use
% 
%   |              |              |              |
%   |    |    |    |    |    |    |    |    |    |   <-- fine mesh
%   |              |              |              |
%  1/2            3/2            5/2            7/2  <-- coarse mesh
%
%  At the coarse mesh boundaries, we need to compute the partial currents
%  going right (Jp) and left (Jm).  Recall, 
%       Jp(i) = int( psi(i,mu), mu =  0..1 )
%       Jm(i) = int( psi(i,mu), mu = -1..0 )   .
%  We also need the removal collision rate in each coarse region.  For a
%  one group problem, this is just the absorption rate.  Then
%       R(i) = int( int(psi(x,mu),mu=-1..1)*SigR(x), x_{i-1/2}..x_{i+1/2}).
%  We finally need the coarse-mesh integrated fluxes,
%       S(i) = int( int(S(x,mu),mu=-1..1), x_{i-1/2}..x_{i+1/2}).
%  With Jp/m(i), R(i), and S(i), we introduce rebalance factors f(i) that
%  will yield neutron balance at the coarse scale.  They are applied to the
%  angular flux (and associated quantities) as follows (see Eq. 20.28 of
%  the 106 notes)
%                       | f(j)*psi(i)^{m+1/2},   x(j-1/2) < x(i) < x(j+1/2)
%       psi(i)^{m+1} = <  f(j-1)*psi(i)^{m+1/2}, x(i) = x(j+1/2)
%                       | f(j+1)*psi(i)^{m+1/2), x(i) = x(j-1/2)
% Substituting these new values into the neutron balance equation
%     Jnet(i+1/2)-Jnet(i-1/2)+R(i)=S(i)
% we get a tridiagonal system for f.  The angular flux is then updated.

Jp = zeros( length(in.xcm), 1 ); % coarse mesh edge rightward p. current
Jm = zeros( length(in.xcm), 1 ); % coarse mesh edge leftward p. current
R  = zeros( length(in.xfm), 1 ); % coarse mesh removal rate
S  = zeros( length(in.xfm), 1 ); % coarse mesh source (no in-scatter)

s = zeros(sum(in.xfm),1);  % the angle-integrated source
for i = 1:length(s)
    s(i) = sum( Q(i,:,1)'.*w(:) );
end

% compute the coarse mesh quantities
Jp(1) = partcur( 1, psi, 1, mu, w, in );
Jm(1) = partcur(-1, psi, 1, mu, w, in );
for i = 1:length(in.xfm)  
    idx1 = 1 + sum( in.xfm(1:(i-1)) );   % lower index
    idx2 = sum( in.xfm(1:(i  )) );       % upper index     
    Jp(i+1) = partcur( 1, psi, idx2+1, mu, w, in );
    Jm(i+1) = partcur(-1, psi, idx2+1, mu, w, in ); 
    R(i) = sum( dx(idx1:idx2)'.*phi(idx1:idx2,1)*in.data( in.mt(i), 2 ) );
    S(i) = sum( dx(idx1:idx2)'.*s(idx1:idx2,1) );
end

% build the tridiagonal system
A = diag(R);
A = A + diag(Jp(2:end)) + diag(Jm(1:end-1)) ...
      - diag(Jp(2:end-1),-1) - diag(Jm(2:end-1),1);
% solve for the rebalance factors
f = A\S;

% update:  
%    within meshes    
%           psi(ii,m) = psi(ii,m)*f(i),    ii E i
%    add coarse mesh edges
%           psi(ii,m) = psi(ii,m)*f(i),    ii E S(i,i')  and mu > 0
%           psi(ii,m) = psi(ii,m)*f(i'),   ii E (Si,i')  and mu < 0
for i = 1:length(in.xfm)  
    % within cell indices
    idx1 = 1 + sum( in.xfm(1:(i-1)) ); % lower index
    idx2 = sum( in.xfm(1:(i  )) );     % upper index   
    psi(idx1+1:idx2,:,1) = psi(idx1+1:idx2,:,1)  * f(i);
    phi(idx1:idx2,1)  = phi(idx1:idx2,1)  * f(i);  %<--not updating all phi's!!
end
idxa1 = in.ord/2+1:in.ord;   % positive
idxa2 = 1:in.ord/2;          % negative
for i = 1:length(in.xfm)+1
    idx  = 1 + sum( in.xfm(1:(i-1)) );
    if ( i == 1 )
        f1 = 0; f2 = f(1);
    elseif ( i == length(in.xfm)+1 )
        f1 = f(end); f2 = 0;
    else
        f1 = f(i-1); f2=f(i);
    end
    psi(idx,idxa1,1) = psi(idx,idxa1,1)*f1; %  -->
    psi(idx,idxa2,1) = psi(idx,idxa2,1)*f2; % <--
end

end

function J = partcur( flag, psi, i, mu, w, in )
    if ( flag == 1 )
        idxs = in.ord/2+1:in.ord;
    else
        idxs = 1:in.ord/2;
    end
    J = sum( abs(mu(idxs)) .* w(idxs) .* psi(i,idxs)' );
end
